# coding: utf-8

import pymongo
import os
import sys
import hashlib
from PIL import Image
from colorHash import color_hash
from os.path import basename
import redis


def getImages(dirname):
    if not os.path.isdir(dirname):
        print("Make sure %s is directory! " % dirname)
        sys.exit(1)
    for root, dirname, fnames in os.walk(dirname):
        for fname in fnames:
            # images.append(os.path.join(root, fname))
            yield os.path.join(root, fname)


def mongoConnection(host="192.168.30.251", dbname="imgSimilar", collname="colorHash"):
    mongoclient = pymongo.MongoClient(host=host)
    mongoclient.admin.authenticate("dmteam", "dmteam")
    coll = mongoclient[dbname][collname]
    return coll


def redisConnection(host="192.168.30.251"):
    rc = redis.Redis(host=host)
    return rc


global DAO
DAO = mongoConnection()  # mongodb
# DAO = redisConnection()  # redis


def md5(fn):
    return hashlib.md5(open(fn).read()).hexdigest()


def buildDoc(fn, hash=None):
    if hash is None:
        hash = color_hash(Image.open(fn))
    doc = {"_id": md5(fn), "hash": str(hash), "name": basename(fn)}
    return doc


def insertDB(doc):
    try:
        DAO.insert(doc)
        return doc
    except Exception as e:
        # print e
        pass
    return None

if __name__ == "__main__":
    # dirname = "/storage/samba/share/workspace/xi_an_test/x"
    dirname = "/storage/samba/share/workspace/similarDB"
    images = getImages(dirname)
    for i in images:
        try:
            doc = buildDoc(i)
            insertDB(doc)
        except: # IOError etc.
            pass

